package dkbta.processing;

import static dkbta.ontology.elements.ElementType.EVENT;
import static dkbta.ontology.elements.ElementType.PRIMITIVE;
import static dkbta.ontology.elements.ElementType.STAT;
import static dkbta.util.Utils.enumForName;
import static dkbta.util.Utils.isEmpty;
import static dkbta.util.Utils.readDOMDocFromXML;
import static dkbta.util.Utils.getDocument;
import static java.lang.Math.max;
import static java.lang.Math.min;

import java.io.File;
import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import dkbta.ontology.Ontology;
import dkbta.ontology.elements.ElementDef;
import dkbta.ontology.elements.ElementType;
import dkbta.ontology.loading.OntologyParser;
import dkbta.processing.data.DB;
import dkbta.processing.data.DataDispatcher;
import dkbta.processing.data.SendTo;
import dkbta.processing.tasks.LocalTask;
import dkbta.processing.tasks.MultiplePatientTask;
import dkbta.processing.tasks.PrimaryRemoteTask;
import dkbta.util.Pair;
import dkbta.util.Utils;

/**
 * The loader for unit plans
 * 
 * @author Uri Kanonov
 */
public class PlanLoader{
	/** The parsed ontology */
	private final Ontology _ontology;

	/** A reference to the data dispatcher */
	private final DataDispatcher _dataDispatcher;

	/** A reference to the work queue */
	private final WorkQueue _workQueue;

	/** A reference to the patient container */
	private final Patients _patients;

	/** The data and control ports per unit (key is unit id) */
	private final Map<Integer, UnitPorts> _unitPorts;

	/**
	 * The elements that can be computed for single patients. This excludes raw data
	 * (Primitives and Events) and multi-patient elements such as Stat
	 */
	private final Set<ElementType> _singlePatientComputableElements = EnumSet
			.complementOf(EnumSet.of(PRIMITIVE, EVENT, STAT));

	/** Database information */
	private final DB _db;

	/** The address of the coordinator */
	private final InetSocketAddress _coordinatorAddress;

	/**
	 * A flag indicating whether the raw data (Primitives and Events) for the
	 * initialization of the patients should come from the database
	 */
	private boolean _rawInputFromDB;

	/**
	 * The name of the file containing the raw data (Primitives and Events) for the
	 * initialization of the patients
	 */
	private String _rawInputFromFile;

	private final StringBuilder _dbQuery;

	/**
	 * Initializes a plan loader and automatically parses the plan referenced by
	 * <tt>planFile</tt>
	 * 
	 * @param planFile The plan to be parsed
	 * @param patients Reference to the patients container
	 * @param dd Reference to the data dispatcher
	 * @param wq Reference to the work queue
	 * @throws Exception If an error occurs during the parsing
	 */
	public PlanLoader(File planFile, Patients patients, DataDispatcher dd, WorkQueue wq)
			throws Exception{
		_patients = patients;
		_dataDispatcher = dd;
		_workQueue = wq;

		// Loading the XML into a DOM document
		Element plan = readDOMDocFromXML(planFile).getDocumentElement();

		_ontology = parseOntology(plan);
		_coordinatorAddress = parseCoordinatorAddress(plan);
		_db = parseDbInfo(plan);
		if (_db != null){
			_dataDispatcher.setDB(_db);
		}

		_dbQuery = initializeDbQuery();
		_unitPorts = parseUnits(plan);
		parsePatients(plan);
		parseStatisticalAbstractions(plan);
	}

	public PlanLoader(String planStr, Patients patients, DataDispatcher dd, WorkQueue wq)
	throws Exception{
		_patients = patients;
		_dataDispatcher = dd;
		_workQueue = wq;

		// Loading the XML into a DOM document
		Element plan = getDocument(planStr);
		
		_ontology = parseOntology(plan);
		_coordinatorAddress = parseCoordinatorAddress(plan);
		_db = parseDbInfo(plan);
		if (_db != null){
			_dataDispatcher.setDB(_db);
		}

		_dbQuery = initializeDbQuery();
		_unitPorts = parseUnits(plan);
		parsePatients(plan);
		parseStatisticalAbstractions(plan);
	}
	
	/**
	 * Parses the ontology referenced by the 'Ontology' element's 'file' attribute
	 * 
	 * @param plan The plan's parsed DOM document
	 * @return The parsed ontology
	 * @throws Exception If the ontology could not be parsed
	 */
	private Ontology parseOntology(Element plan) throws Exception{
		Element ontologyElement = getXmlElement(plan, "Ontology");
		if (ontologyElement == null){
			throw new PlanException("No 'Ontology' element in the plan");
		}
		File ontologyFile = new File(ontologyElement.getAttribute("file"));
		OntologyParser ontLoader = new OntologyParser();
		return ontLoader.load(ontologyFile);
	}

	/**
	 * Parses the database info
	 * 
	 * @param plan The plan's parsed DOM document
	 * @return The parsed database info or null if there was no database info in the plan
	 * @throws PlanException If the info was malformed
	 */
	private DB parseDbInfo(Element plan) throws PlanException{
		Element dbInfoElement = getXmlElement(plan, "DataBaseInfo");
		if (dbInfoElement == null || !dbInfoElement.hasAttributes()){
			return null; // We won't be able to send info the DB
		}

		String dbPassword = dbInfoElement.getAttribute("dbPassword");
		String dbUsername = dbInfoElement.getAttribute("dbUsername");
		String dbURL = dbInfoElement.getAttribute("dbURL");
		String dbDriver = dbInfoElement.getAttribute("dbDriver");

		if (isEmpty(dbPassword) || isEmpty(dbUsername) || isEmpty(dbURL)
				|| isEmpty(dbDriver)){
			throw new PlanException("Invalid DB Info");
		}

		String inputTableName = dbInfoElement.getAttribute("inputTableName");
		String outputTableName = dbInfoElement.getAttribute("outputTableName");

		return new DB(dbURL, dbDriver, dbUsername, dbPassword, inputTableName,
				outputTableName);
	}

	/**
	 * Parses the information regarding each of the patients defined in the plan (local
	 * and remote elements)
	 * 
	 * @param plan The plan's parsed DOM document
	 * @throws PlanException If the plan is malformed
	 */
	private void parsePatients(Element plan) throws Exception{
		Element patientsNode = getXmlElement(plan, "Patients");
		if (patientsNode == null || !patientsNode.hasAttributes()){
			throw new PlanException(
					"There is no patients node with attributes indicating raw input source");
		}

		NodeList patientNodes = plan.getElementsByTagName("Patient");
		int length = patientNodes.getLength();
		if (length == 0){
			throw new PlanException("At least one patient must be specified");
		}
		for (int i = 0; i < length; ++i){
			Element patient = (Element)patientNodes.item(i);

			// Parsing the patient ids
			String idStr = patient.getAttribute("id");
			Pair<Integer, Integer>[] patientIdRanges = parsePatientIds(idStr);

			// Parsing the remote elements
			Set<ElementDef> remoteElements;
			try{
				remoteElements = parseRemoteElements(patient);
			}catch(PlanException e){
				throw new PlanException("Unable to parse remote elements for patients: "
						+ idStr, e);
			}

			// Parsing the local elements
			Map<ElementDef, SendTo> localElements;
			try{
				localElements = parseLocalElements(patient, remoteElements);
			}catch(PlanException e){
				throw new PlanException("Unable to parse local elements for patients: "
						+ idStr, e);
			}

			if (localElements.isEmpty() && remoteElements.isEmpty()){
				throw new PlanException("No local/remote elements "
						+ "have been specified for patients: " + idStr);
			}

			// For all patients
			for (Pair<Integer, Integer> idRange : patientIdRanges){
				int start = idRange.getFirst();
				int end = idRange.getSecond();
				for (int id = start; id <= end; ++id){
					// Creating the patient
					Patient p = new Patient(id, _workQueue, _dataDispatcher,
							remoteElements, localElements);
					_patients.addPatient(p);

					// Adding all of the patient's local elements to the work queue
					for (ElementDef ed : localElements.keySet()){
						_workQueue.newTask(new LocalTask(ed, true, p), true);
					}

					// Wrapping all of the remote elements with dummy primary tasks
					// and adding them to the work queue
					for (ElementDef ed : remoteElements){
						_workQueue.newTask(new PrimaryRemoteTask(ed, p), true);
					}
				}
			}

			addPatientsToQuery(patientIdRanges, localElements.keySet(), i == length -1);
		}

		parsePatientsInitializationMode(patientsNode);
		
		finalizeQuery();
	}

	/**
	 * Parses an id-range
	 * 
	 * @param idRangesStr The id ranges in textual form
	 * @return The parsed id ranges
	 * @throws PlanException If the textual id ranges are malformed
	 */
	private Pair<Integer, Integer>[] parsePatientIds(String idRangesStr)
			throws PlanException{
		String[] idRanges = idRangesStr.trim().split(" *, *");
		if (idRanges.length == 0){
			throw new PlanException("No patient ids specified");
		}
		@SuppressWarnings("unchecked")
		Pair<Integer, Integer>[] ids = new Pair[idRanges.length];
		for (int j = 0; j < idRanges.length; ++j){
			String[] idRange = idRanges[j].split(" *- *");
			int rangeLen = idRange.length;
			if (rangeLen > 2){
				throw new PlanException("Invalid patient ids: " + idRangesStr);
			}
			try{
				int id1 = Integer.valueOf(idRange[0]);
				int id2 = rangeLen == 2 ? Integer.valueOf(idRange[1]) : id1;
				ids[j] = new Pair<Integer, Integer>(min(id1, id2), max(id1, id2));
			}catch(NumberFormatException e){
				throw new PlanException("Invalid patient ids: " + idRangesStr);
			}
		}
		return ids;
	}

	/**
	 * Parses the attributes of the 'Patients' node that say how to initialize the raw
	 * data (from file or DB)
	 * 
	 * @param patientsNode The 'Patients' node
	 * @throws PlanException If the data is malformed
	 */
	private void parsePatientsInitializationMode(Element patientsNode)
			throws PlanException{
		_rawInputFromDB = Boolean.parseBoolean(patientsNode
				.getAttribute("rawInputFromDB"));
		_rawInputFromFile = patientsNode.getAttribute("rawInputFromFile");
		if (_rawInputFromDB){
			if (_db == null){
				throw new PlanException(
						"Cannot obtain raw input from DB without DB info, fill the DataBaseInfo element");
			}
			if (!isEmpty(_rawInputFromFile)){
				throw new PlanException(
						"Must specify only one raw data source in the patients element (not both)!");
			}
		}else{
			if (isEmpty(_rawInputFromFile)){
				throw new PlanException(
						"Must specify a raw data source in the patients element!");
			}
		}
	}

	/**
	 * Parses the local elements defined for the given patient
	 * 
	 * @param patient The patient's XML node
	 * @param remoteElements The patient's remote elements
	 * @return The parsed local elements (with send-to information)
	 * @throws PlanException If the plan is malformed (e.g. an element is defined both as
	 *         local and remote)
	 */
	private Map<ElementDef, SendTo> parseLocalElements(Element patient,
		Set<ElementDef> remoteElements) throws PlanException{
		Map<ElementDef, SendTo> localElements = new HashMap<ElementDef, SendTo>();
		Element localsNode = getXmlElement(patient, "Local");
		if (localsNode == null){
			return localElements;
		}
		NodeList localNodes = localsNode.getChildNodes();
		int length = localNodes.getLength();
		for (int i = 0; i < length; ++i){
			Node node = localNodes.item(i);
			if (node.getNodeType() != Node.ELEMENT_NODE){
				continue;
			}
			Element localElement = (Element)localNodes.item(i);

			// Parsing the send-to information
			String sendToUnits = localElement.getAttribute("sendToUnits");
			Collection<InetSocketAddress> units = null;
			if (!isEmpty(sendToUnits)){
				String[] unitsArr = sendToUnits.split(" *, *");
				units = new ArrayList<InetSocketAddress>(unitsArr.length);
				for (String unitIdStr : unitsArr){
					InetSocketAddress unitDataPort = null;
					try{
						unitDataPort = _unitPorts.get(Integer.valueOf(unitIdStr)).dataPort;
					}catch(NumberFormatException e){}
					if (unitDataPort == null){
						throw new PlanException(
								"Invalid/Unknown unit-id in 'sendToUnits': " + unitIdStr);
					}
					units.add(unitDataPort);
				}
			}
			boolean sendToDB = Boolean
					.parseBoolean(localElement.getAttribute("sendToDB"));
			SendTo sendTo = (sendToDB || units != null) ? new SendTo(units, sendToDB)
					: null;

			// Parsing the ontology elements
			Collection<? extends ElementDef> elements = parseElement(localElement,
				_singlePatientComputableElements);
			for (ElementDef ed : elements){
				if (remoteElements.contains(ed)){
					throw new PlanException("The " + ed
							+ " cannot be both remote and local");
				}else if (localElements.containsKey(ed)){
					throw new PlanException("The " + ed
							+ " is specified twice in the local element list");
				}
				localElements.put(ed, sendTo);
			}
		}

		return localElements;
	}

	/**
	 * Parses the remote elements from the given patient node
	 * 
	 * @param patient The patient's XML node
	 * @return The patient's remote elements
	 * @throws PlanException If the plan is malformed
	 */
	@SuppressWarnings("unchecked")
	private Set<ElementDef> parseRemoteElements(Element patient) throws PlanException{
		Element localsNode = getXmlElement(patient, "Remote");
		if (localsNode == null){
			return Collections.EMPTY_SET;
		}

		NodeList remoteNodes = localsNode.getChildNodes();
		int length = remoteNodes.getLength();
		Set<ElementDef> remoteElements = new HashSet<ElementDef>(length);
		for (int i = 0; i < length; ++i){
			Node node = remoteNodes.item(i);
			if (node.getNodeType() != Node.ELEMENT_NODE){
				continue;
			}
			Element remoteElement = (Element)remoteNodes.item(i);

			// Parsing the ontology element
			remoteElements.addAll(parseElement(remoteElement,
				_singlePatientComputableElements));
		}

		return Collections.unmodifiableSet(remoteElements);
	}

	/**
	 * Parses an XML element defining either a single ontology element (if the 'name'
	 * attribute is defined and valid) or all of the elements of a certain type (if there
	 * is no 'name' attribute)
	 * 
	 * @param element The XML element defining the ontology element(s)
	 * @param validElementTypes Which element types are considered valid
	 * @return The defined element(s). If it was a single element, it is wrapped in a
	 *         collection
	 * @throws PlanException If the element definition is invalid, the specified element
	 *         doesn't exist in the ontology or the type is invalid
	 */
	private Collection<? extends ElementDef> parseElement(Element element,
		Set<ElementType> validElementTypes) throws PlanException{
		ElementType et = enumForName(ElementType.class, element.getTagName());
		if (!validElementTypes.contains(et)){
			throw new PlanException("Invalid element type: " + et);
		}

		if (element.hasAttribute("name")){
			String name = element.getAttribute("name");
			if (Utils.isEmpty(name)){
				throw new PlanException("Invalid element name: " + name);
			}
			ElementDef ed = _ontology.getElement(et, name);
			if (ed == null){
				throw new PlanException("Element not defined in ontology: " + et + " - "
						+ name);
			}
			return Arrays.asList(ed);
		}else{
			return _ontology.getAllElements(et);
		}
	}

	/**
	 * Parses the statistical abstractions defined in the plan
	 * 
	 * @param plan The plan's parsed DOM document
	 * @throws PlanException If the plan is malformed
	 */
	private void parseStatisticalAbstractions(Element plan) throws Exception{
		Element statsticalNode = getXmlElement(plan, "Statistical");
		if (statsticalNode == null){
			return; // No statistical abstractions
		}

		Set<ElementType> statAbstractionTypes = EnumSet.of(ElementType.STAT);

		Map<ElementDef, SendTo> statisticalAbstractions = new HashMap<ElementDef, SendTo>();
		NodeList statNodes = statsticalNode.getChildNodes();
		int length = statNodes.getLength();
		for (int i = 0; i < length; ++i){
			Node node = statNodes.item(i);
			if (node.getNodeType() != Node.ELEMENT_NODE){
				continue;
			}
			Element statElement = (Element)statNodes.item(i);

			// Parsing the statistical abstraction ontology elements
			Collection<? extends ElementDef> stats = parseElement(statElement,
				statAbstractionTypes);

			// Parsing the patient ids
			String idStr = statElement.getAttribute("patients");
			Pair<Integer, Integer>[] ids = parsePatientIds(idStr);

			// Checking that all patients have been defined in the 'Patients' section
			try{
				Iterator<Patient> iter = _patients.rangeIterator(ids);
				while (iter.hasNext()){
					Patient p = iter.next();

					p.assertComputationOfPrerequisites(stats);
				}
			}catch(UnknownPatientException e){
				throw new PlanException(
						"Reference in a statistical element to undefined patient: "
								+ e.getPatientId());
			}

			// Parsing the send-to-db flag
			boolean sendToDB = Boolean.parseBoolean(statElement.getAttribute("sendToDB"));
			SendTo sendTo = (sendToDB) ? new SendTo(null, sendToDB) : null;

			// Adding tasks to compute each of the statistical abstractions
			for (ElementDef element : stats){
				_workQueue
						.newTask(new MultiplePatientTask(element, _patients, ids), true);
				statisticalAbstractions.put(element, sendTo);
			}

			// Adding the statistical abstractions along with their send-to info to the
			// Patients object
			_patients.setSendTos(statisticalAbstractions);
		}
	}

	/**
	 * Parses the 'Units' node containing the unit ids, data and control ports
	 * 
	 * @param plan The plan node
	 * @return The parsed unit ports
	 * @throws PlanException If the plan is invalid
	 */
	public static Map<Integer, UnitPorts> parseUnits(Element plan) throws PlanException{
		Element unitsElement = getXmlElement(plan, "Units");
		if (unitsElement == null){
			throw new PlanException("There is no 'Units' element");
		}

		NodeList unitNodes = unitsElement.getElementsByTagName("Unit");
		int length = unitNodes.getLength();
		if (length == 0){
			throw new PlanException(
					"There are no 'Unit' elements inside the 'Units' element");
		}

		Map<Integer, UnitPorts> units = new HashMap<Integer, UnitPorts>(length);
		for (int i = 0; i < length; ++i){
			Element unit = (Element)unitNodes.item(i);

			// Parsing the unit id
			String idStr = unit.getAttribute("id");
			int id;
			try{
				id = Integer.valueOf(idStr);
			}catch(NumberFormatException e){
				throw new PlanException("Invalid unit id: " + idStr);
			}

			// Parsing the unit's ip/hostname
			String ip = unit.getAttribute("ip");
			if (isEmpty(ip)){
				throw new PlanException("Invalid ip '" + ip + "' of unit #" + id);
			}

			// Parsing the unit's data port
			String dataPortStr = unit.getAttribute("dataPort");
			int dataPort;
			try{
				dataPort = Integer.valueOf(dataPortStr);
			}catch(NumberFormatException e){
				throw new PlanException("Invalid data port '" + dataPortStr
						+ "' of unit #" + id);
			}

			// Parsing the unit's data port
			String controlPortStr = unit.getAttribute("controlPort");
			int controlPort;
			try{
				controlPort = Integer.valueOf(controlPortStr);
			}catch(NumberFormatException e){
				throw new PlanException("Invalid control port '" + controlPortStr
						+ "' of unit #" + id);
			}
			units.put(id, new UnitPorts(new InetSocketAddress(ip, dataPort),
					new InetSocketAddress(ip, controlPort)));
		}

		return units;
	}

	/**
	 * Parses the coordinator's address form the 'Coordinator' node
	 * 
	 * @param plan The plan node
	 * @return The coordinator's address
	 * @throws PlanException If the plan is invalid
	 */
	public static InetSocketAddress parseCoordinatorAddress(Element plan)
			throws PlanException{
		Element coordinatorElement = getXmlElement(plan, "Coordinator");
		if (coordinatorElement == null){
			throw new PlanException("There is no 'Coordinator' element");
		}

		// Parsing the coordinator's ip/hostname
		String ip = coordinatorElement.getAttribute("ip");
		if (isEmpty(ip)){
			throw new PlanException("Invalid ip '" + ip + "' of the coordinator");
		}

		// Parsing the coordinator's port
		String portStr = coordinatorElement.getAttribute("port");
		int port;
		try{
			port = Integer.valueOf(portStr);
		}catch(NumberFormatException e){
			throw new PlanException("Invalid port '" + portStr + "' of the coordinator");
		}
		return new InetSocketAddress(ip, port);
	}
	
	private static Element getXmlElement(Element parent, String tagName)
			throws PlanException{
		NodeList nodes = parent.getElementsByTagName(tagName);
		if (nodes.getLength() == 0){
			return null;
		}
		if (nodes.getLength() != 1){
			throw new PlanException("Multile '" + tagName + "' elements");
		}

		return (Element)nodes.item(0);
	}

	/**
	 * Retrieves the parsed ontology
	 * 
	 * @return The parsed ontology
	 */
	public Ontology getOntology(){
		return _ontology;
	}

	/**
	 * Retrieves the unit ports for the given unit
	 * 
	 * @param uid The id of the unit
	 * @return The data and control ports of the given unit
	 * @throws IllegalStateException If the unit was not defined in the plan
	 */
	public UnitPorts getPorts(int uid){
		UnitPorts ports = _unitPorts.get(uid);
		if (ports == null){
			throw new IllegalStateException("This unit itself (" + uid
					+ ") is not defined in the plan!");
		}
		return ports;
	}

	/**
	 * Retrieves the address (hostname, port) of the coordinator
	 * 
	 * @return The address (hostname, port) of the coordinator
	 */
	public InetSocketAddress getCoordinatorAddress(){
		return _coordinatorAddress;
	}

	/**
	 * Performs the initialization of the patients according to the attributes from the
	 * 'Patients' node in the plan
	 * 
	 * @throws Exception If an error occurs during the initialization
	 */
	public void initializePatients() throws Exception{
		if (_rawInputFromDB){
			_patients.initializeFromDB(_ontology, _db, _dbQuery.toString());
		}else{
			_patients.initializeFromFile(_ontology, _rawInputFromFile);
		}
	}

	private StringBuilder initializeDbQuery(){
		if (_db == null){
			return null;
		}
		StringBuilder query = new StringBuilder();

		query.append("SELECT SubjectID, FeatureName, StartTime, EndTime, FeatureValue\n");
		query.append("FROM ").append(_db.getInputTableName()).append("\n");
		query.append("WHERE ");

		return query;
	}

	private void addPatientsToQuery(Pair<Integer, Integer>[] patientIdRanges, Set<ElementDef> set, boolean areLastPatients){
		Patient p = _patients.getPatient(patientIdRanges[0].getFirst());
		Set<ElementDef> requiredPrimsAndEvents = p.getRequiredPrimsAndEvents();
		for (int i = 0; i < patientIdRanges.length; ++i){
			Pair<Integer, Integer> idRange = patientIdRanges[i];
			int min = idRange.getFirst();
			int max = idRange.getSecond();
			if (min == max){
				_dbQuery.append("\n(SubjectID = ").append(min);
			}else{
				_dbQuery.append("\n((SubjectID >= ").append(min).append(
					" AND SubjectID <= ").append(max).append(")");
			}
			if (i < patientIdRanges.length - 1){
				_dbQuery.append(" OR ");
			}
		}

		_dbQuery.append(" AND FeatureName in (");
		appendElementNames(requiredPrimsAndEvents, _dbQuery).append("))");
		if (!areLastPatients){
			_dbQuery.append("\nOR ");
		}
	}
	
	private StringBuilder appendElementNames(Collection<ElementDef> eds, StringBuilder sb){
		Iterator<ElementDef> iter = eds.iterator();
		while (iter.hasNext()){
			ElementDef ed = iter.next();
			sb.append("'").append(ed.getName()).append("'");
			if (iter.hasNext()){
				sb.append(",");
			}
		}
		return sb;
	}

	private void finalizeQuery(){
		_dbQuery.append("\nORDER BY SubjectID ASC, FeatureName ASC, StartTime ASC");
	}
	
	public StringBuilder getDbQuery(){
		return _dbQuery;
	}
}
